import baostock as bs
import pandas as pd
from typing import List, Dict, Any
import time
from datetime import datetime

from .utils.config import get_settings, ensure_data_dirs
from .utils.logger import get_logger
from pathlib import Path

logger = get_logger("AShareList")


class AShareListFetcher:
    """A股股票列表获取器"""

    def __init__(self):
        self.settings = get_settings()
        self.logger = logger
        self.bs = None
        self.is_logged_in = False

    def __del__(self):
        """析构函数，确保登出"""
        if self.is_logged_in:
            self.logout()

    def login(self) -> bool:
        """登录baostock"""
        try:
            if self.is_logged_in:
                return True

            self.logger.info("登录baostock...")
            result = bs.login()
            if result.error_code != '0':
                raise Exception(f"登录失败: {result.error_msg}")

            self.is_logged_in = True
            self.bs = bs
            self.logger.info("baostock登录成功")
            return True

        except Exception as e:
            self.logger.error(f"baostock登录失败: {e}")
            return False

    def logout(self):
        """登出baostock"""
        try:
            if self.is_logged_in and self.bs:
                self.bs.logout()
                self.is_logged_in = False
                self.bs = None
                self.logger.info("baostock已登出")
        except Exception as e:
            self.logger.error(f"baostock登出失败: {e}")

    def get_all_stocks(self) -> pd.DataFrame:
        """获取所有A股股票列表"""
        try:
            if not self.login():
                raise Exception("登录失败")

            self.logger.info("获取所有A股股票列表...")

            # 获取证券信息
            rs = bs.query_stock_basic()
            if rs.error_code != '0':
                raise Exception(f"获取股票列表失败: {rs.error_msg}")

            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())

            if not data_list:
                raise Exception("未获取到股票列表")

            # 创建DataFrame
            df = pd.DataFrame(data_list, columns=rs.fields)

            # 数据清洗
            df = self._clean_stock_list(df)

            self.logger.info(f"成功获取 {len(df)} 只A股股票")
            return df

        except Exception as e:
            self.logger.error(f"获取A股股票列表失败: {e}")
            raise

    def _clean_stock_list(self, df: pd.DataFrame) -> pd.DataFrame:
        """清洗股票列表数据"""
        try:
            # 删除空值
            df = df.replace('', None).dropna()

            if df.empty:
                return df

            # 筛选股票类型（排除指数、基金等）
            # type_name字段：股票、指数、基金、债券等
            if 'type_name' in df.columns:
                df = df[df['type_name'] == '股票']

            # 筛选状态为上市的股票
            if 'status_name' in df.columns:
                df = df[df['status_name'] == '上市']

            # 删除退市股票
            if 'out_date' in df.columns:
                df = df[df['out_date'].isna()]

            # 添加股票代码后缀用于显示
            if 'code' in df.columns:
                df['symbol'] = df['code'].str.replace('sh.', '').str.replace('sz.', '').str.replace('bj.', '')

            # 添加交易所信息
            if 'code' in df.columns:
                df['exchange'] = df['code'].str[:2].map({
                    'sh': '上海证券交易所',
                    'sz': '深圳证券交易所',
                    'bj': '北京证券交易所'
                }).fillna('未知')

            # 添加市场类型
            df['market_type'] = df['symbol'].apply(self._get_market_type)

            # 重命名列
            column_mapping = {
                'code': 'bs_code',
                'code_name': 'name',
                'ipo_date': 'ipo_date',
                'out_date': 'delist_date',
                'type_name': 'type',
                'status_name': 'status'
            }

            df = df.rename(columns=column_mapping)

            # 选择需要的列
            columns_to_keep = [
                'symbol', 'bs_code', 'name', 'exchange', 'market_type',
                'type', 'status', 'ipo_date', 'delist_date'
            ]
            available_columns = [col for col in columns_to_keep if col in df.columns]
            df = df[available_columns]

            # 按股票代码排序
            df = df.sort_values('symbol').reset_index(drop=True)

            return df

        except Exception as e:
            self.logger.error(f"股票列表数据清洗失败: {e}")
            return df

    def _get_market_type(self, symbol: str) -> str:
        """根据股票代码判断市场类型"""
        if not symbol:
            return '未知'

        if symbol.startswith('688'):
            return '科创板'
        elif symbol.startswith('600') or symbol.startswith('601') or symbol.startswith('602') or symbol.startswith('603') or symbol.startswith('605'):
            return '上海主板'
        elif symbol.startswith('000') or symbol.startswith('001') or symbol.startswith('002'):
            return '深圳主板'
        elif symbol.startswith('300'):
            return '创业板'
        elif symbol.startswith('43') or symbol.startswith('83') or symbol.startswith('87'):
            return '新三板'
        elif symbol.startswith('8'):
            return '北交所'
        else:
            return '其他'

    def get_stocks_by_market(self, market_type: str) -> List[str]:
        """根据市场类型获取股票代码列表"""
        try:
            df = self.get_all_stocks()
            market_df = df[df['market_type'] == market_type]
            return market_df['symbol'].tolist()
        except Exception as e:
            self.logger.error(f"获取{market_type}股票列表失败: {e}")
            return []

    def get_active_stocks(self) -> List[str]:
        """获取所有活跃交易的股票代码"""
        try:
            df = self.get_all_stocks()
            active_df = df[df['status'] == '上市']
            return active_df['symbol'].tolist()
        except Exception as e:
            self.logger.error(f"获取活跃股票列表失败: {e}")
            return []

    def get_stock_info(self, symbol: str) -> Dict[str, Any]:
        """获取单个股票的详细信息"""
        try:
            df = self.get_all_stocks()
            stock_df = df[df['symbol'] == symbol]
            if stock_df.empty:
                return {}
            return stock_df.iloc[0].to_dict()
        except Exception as e:
            self.logger.error(f"获取股票{symbol}信息失败: {e}")
            return {}

    def save_stock_list(self, file_path: str = None) -> bool:
        """保存股票列表到CSV文件"""
        try:
            if file_path is None:
                from .utils.config import ensure_data_dirs
                dirs = ensure_data_dirs()
                file_path = dirs['data_dir'] / 'metadata' / 'ashare_list.csv'

            df = self.get_all_stocks()
            df.to_csv(file_path, index=False, encoding='utf-8')

            self.logger.info(f"股票列表已保存到: {file_path}")
            return True

        except Exception as e:
            self.logger.error(f"保存股票列表失败: {e}")
            return False

    def load_stock_list(self, file_path: str = None) -> pd.DataFrame:
        """从CSV文件加载股票列表"""
        try:
            if file_path is None:
                from .utils.config import ensure_data_dirs
                dirs = ensure_data_dirs()
                file_path = dirs['data_dir'] / 'metadata' / 'ashare_list.csv'

            if not Path(file_path).exists():
                self.logger.warning(f"股票列表文件不存在: {file_path}")
                return pd.DataFrame()

            df = pd.read_csv(file_path, encoding='utf-8')
            self.logger.info(f"从文件加载股票列表: {len(df)} 只股票")
            return df

        except Exception as e:
            self.logger.error(f"加载股票列表失败: {e}")
            return pd.DataFrame()

    def get_market_summary(self) -> Dict[str, int]:
        """获取市场统计信息"""
        try:
            df = self.get_all_stocks()

            summary = {
                'total': len(df),
                'sh_main': len(df[df['market_type'] == '上海主板']),
                'sz_main': len(df[df['market_type'] == '深圳主板']),
                'star': len(df[df['market_type'] == '科创板']),
                'chi_next': len(df[df['market_type'] == '创业板']),
                'new_obb': len(df[df['market_type'] == '新三板']),
                'beijing': len(df[df['market_type'] == '北交所']),
                'active': len(df[df['status'] == '上市'])
            }

            return summary

        except Exception as e:
            self.logger.error(f"获取市场统计失败: {e}")
            return {}

    def refresh_stock_list(self) -> bool:
        """刷新股票列表"""
        try:
            self.logger.info("开始刷新股票列表...")
            success = self.save_stock_list()
            if success:
                summary = self.get_market_summary()
                self.logger.info(f"股票列表刷新成功: {summary}")
            return success
        except Exception as e:
            self.logger.error(f"刷新股票列表失败: {e}")
            return False